"use client";
import {
  ReactFlow,
  Node,
  Edge,
  useNodesState,
  useEdgesState,
  Panel,
  useReactFlow,
  Background,
  NodeTypes,
} from "@xyflow/react";
import { useEffect, useMemo } from "react";
import { Button } from "@/components/ui/button";
import { ChevronUp, ChevronDown } from "lucide-react";
import { ImageNode } from "./image-node"; // 导入 ImageNode
import { AgentNode } from "./agent-node"; // 导入 AgentNode

interface ValidationFlowProps {
  rlName: string;
  cvModels: string[];
  inputImages: string[];
  outputImages: string[];
  chosenModel: string;
  onSwitchImageGroup?: (direction: "up" | "down") => void;
}

export default function ValidationFlow({
  rlName,
  cvModels,
  inputImages,
  outputImages,
  chosenModel,
  onSwitchImageGroup,
}: ValidationFlowProps) {
  const [nodes, setNodes, onNodesChange] = useNodesState<Node>([]);
  const [edges, setEdges, onEdgesChange] = useEdgesState<Edge>([]);
  const { fitView } = useReactFlow();

  // 定义节点类型
  const nodeTypes: NodeTypes = useMemo(
    () => ({
      imageNode: ImageNode,
      agentNode: AgentNode, // 添加 AgentNode 到 nodeTypes
    }),
    []
  );

  useEffect(() => {
    const newNodes: Node[] = [];
    const newEdges: Edge[] = [];

    const inputNodeExists = inputImages && inputImages.length > 0;
    const outputNodeExists = outputImages && outputImages.length > 0;

    // 输入图片节点
    if (inputNodeExists) {
      newNodes.push({
        id: "input-images-node",
        type: "imageNode",
        position: { x: 0, y: 0 },
        data: { imageUrls: inputImages },
      });
    }

    // Agent 节点
    newNodes.push({
      id: "agent-node",
      type: "agentNode",
      position: { x: 0, y: 250 }, // 放置在中间
      data: { rlName, cvModels, chosenModel },
    });

    // 输出图片节点
    if (outputNodeExists) {
      newNodes.push({
        id: "output-images-node",
        type: "imageNode",
        position: { x: 0, y: 400 }, // 纵向排列，y轴增加
        data: { imageUrls: outputImages },
      });
    }

    // 创建边
    if (inputNodeExists) {
      newEdges.push({
        id: "input-to-agent-edge",
        source: "input-images-node",
        target: "agent-node",
        animated: true,
      });
    }

    if (outputNodeExists) {
      newEdges.push({
        id: "agent-to-output-edge",
        source: "agent-node",
        target: "output-images-node",
        animated: true,
      });
    }

    setNodes(newNodes);
    setEdges(newEdges);
  }, [
    inputImages,
    outputImages,
    rlName,
    cvModels,
    chosenModel,
    setNodes,
    setEdges,
  ]);

  // 节点和边变化后自动 fitView
  useEffect(() => {
    fitView({ padding: 0.2 });
  }, [nodes, edges, fitView]);

  return (
    <div className="w-full h-full rounded-md border">
      <ReactFlow
        nodes={nodes}
        edges={edges}
        onNodesChange={onNodesChange}
        onEdgesChange={onEdgesChange}
        nodeTypes={nodeTypes}
        fitView
      >
        <Panel position="top-right">
          <div className="flex flex-col gap-2 items-center">
            <Button
              variant="secondary"
              size="icon"
              onClick={() => onSwitchImageGroup?.("up")}
              aria-label="上一个图片组"
            >
              <ChevronUp />
            </Button>
            <Button
              variant="secondary"
              size="icon"
              onClick={() => onSwitchImageGroup?.("down")}
              aria-label="下一个图片组"
            >
              <ChevronDown />
            </Button>
          </div>
        </Panel>
        <Background />
      </ReactFlow>
    </div>
  );
}
